#include <cassert>
#include "libyuv.h"
#include "jni.h"
#include "android/log.h"

#define CLASS_JAVA "top/broncho/yuv/YuvUtils"

#define LOG_TAG "YuvUtils"
#define LOGD(...) __android_log_print(ANDROID_LOG_DEBUG, LOG_TAG, __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__)

#define max(x, y)  (x>y?x:y)
#define min(x, y)  (x<y?x:y)
//#define y(r, g, b)  (((66 * r + 129 * g + 25 * b + 128) >> 8) + 16)
//#define u(r, g, b)  (((-38 * r - 74 * g + 112 * b + 128) >> 8) + 128)
//#define v(r, g, b)  (((112 * r - 94 * g - 18 * b + 128) >> 8) + 128)
#define y(r, g, b)  ((77 * r + 150 * g + 29 * b + 128) >> 8)
#define u(r, g, b)  (((-43 * r - 84 * g + 127 * b + 128) >> 8) + 128)
#define v(r, g, b)  (((127 * r - 106 * g - 21 * b + 128) >> 8) + 128)
#define color(x)  ((unsigned char)((x < 0) ? 0 : ((x > 255) ? 255 : x)))

#define RGBA_YUV420SP   0x00004012
#define BGRA_YUV420SP   0x00004210
#define RGBA_YUV420P    0x00014012
#define BGRA_YUV420P    0x00014210
#define RGB_YUV420SP    0x00003012
#define RGB_YUV420P     0x00013012
#define BGR_YUV420SP    0x00003210
#define BGR_YUV420P     0x00013210

#ifdef __cplusplus
extern "C" {
#endif

static int
(*rgbaToI420Func[])(const uint8_t *, int, uint8_t *, int, uint8_t *, int, uint8_t *, int, int,
                    int) ={
        libyuv::ABGRToI420, libyuv::RGBAToI420, libyuv::ARGBToI420, libyuv::BGRAToI420,
        libyuv::RGB24ToI420, libyuv::RGB565ToI420
};

static int
(*i420ToRgbaFunc[])(const uint8_t *, int, const uint8_t *, int, const uint8_t *, int, uint8_t *,
                    int, int, int) ={
        libyuv::I420ToABGR, libyuv::I420ToRGBA, libyuv::I420ToARGB, libyuv::I420ToBGRA,
        libyuv::I420ToRGB24, libyuv::I420ToRGB565
};

/**
*   type 0-3位表示b的偏移量
*        4-7位表示g的偏移量
*        8-11位表示r的偏移量
*        12-15位表示rgba一个像素所占的byte
*        16-19位表示yuv的类型，0为420sp，1为420p
*/
void rgbaToYuv(int width, int height, const unsigned char *rgb, unsigned char *yuv, int type) {
    const int frameSize = width * height;
    const int yuvType = (type & 0x10000) >> 16;
    const int byteRgba = (type & 0x0F000) >> 12;
    const int rShift = (type & 0x00F00) >> 8;
    const int gShift = (type & 0x000F0) >> 4;
    const int bShift = (type & 0x0000F);
    const int uIndex = 0;
    const int vIndex = yuvType; //yuvType为1表示YUV420p,为0表示420sp

    int yIndex = 0;
    int uvIndex[2] = {frameSize, frameSize + frameSize / 4};

    unsigned char R, G, B, Y, U, V;
    unsigned int index;
    for (int j = 0; j < height; j++) {
        for (int i = 0; i < width; i++) {
            index = j * width + i;

            R = rgb[index * byteRgba + rShift] & 0xFF;
            G = rgb[index * byteRgba + gShift] & 0xFF;
            B = rgb[index * byteRgba + bShift] & 0xFF;

            Y = y(R, G, B);
            U = u(R, G, B);
            V = v(R, G, B);

            yuv[yIndex++] = color(Y);
            if (j % 2 == 0 && index % 2 == 0) {
                yuv[uvIndex[uIndex]++] = color(U);
                yuv[uvIndex[vIndex]++] = color(V);
            }
        }
    }
}

int Jni_ARGBToNv21(JNIEnv *env, jclass clazz, jbyteArray rgba, jbyteArray yuv, jint width,
                   jint height) {
    int rgba_stride = 4 * width;
    LOGD("Jni_ArgbToNv21： rgba_stride=%d", rgba_stride);
    auto ySize = width * height;
    jbyte *rgbaData = env->GetByteArrayElements(rgba, JNI_FALSE);
    jbyte *yuvData = env->GetByteArrayElements(yuv, JNI_FALSE);
    // 传统方式
//    auto *cRgba = (unsigned char *) rgbaData;
//    auto *cYuv = (unsigned char *) yuvData;
//    rgbaToYuv(width, height, cRgba, cYuv, BGRA_YUV420SP);
    // neon 加速
    int ret = libyuv::ABGRToNV21((uint8_t *) rgbaData, rgba_stride, (uint8_t *) yuvData, width,
                                 (uint8_t * )(yuvData + ySize), width, width, height);
    env->ReleaseByteArrayElements(rgba, rgbaData, JNI_OK);
    env->ReleaseByteArrayElements(yuv, yuvData, JNI_OK);
    return ret;
}

int Jni_Nv21ToABGR(JNIEnv *env, jclass clazz, jbyteArray yuv, jbyteArray rgba, jint width,
                   jint height) {
    int rgba_stride = 4 * width;
    LOGD("Jni_Nv21ToABGR： rgba_stride=%d", rgba_stride);
    auto ySize = width * height;
    jbyte *rgbaData = env->GetByteArrayElements(rgba, JNI_FALSE);
    jbyte *yuvData = env->GetByteArrayElements(yuv, JNI_FALSE);
    int ret = libyuv::NV21ToABGR((uint8_t *) yuvData, width,
                                 (uint8_t *) (yuvData + ySize), width, (uint8_t *) rgbaData,
                                 rgba_stride, width, height);
    env->ReleaseByteArrayElements(rgba, rgbaData, JNI_OK);
    env->ReleaseByteArrayElements(yuv, yuvData, JNI_OK);
    return ret;
}

int rgbaToI420(JNIEnv *env, jclass clazz, jbyteArray rgba, jint rgba_stride,
               jbyteArray yuv, jint y_stride, jint u_stride, jint v_stride,
               jint width, jint height,
               int (*func)(const uint8_t *, int, uint8_t *, int, uint8_t *, int, uint8_t *, int,
                           int, int)) {
    size_t ySize = (size_t) (y_stride * height);
    size_t uSize = (size_t) (u_stride * height >> 1);
    jbyte *rgbaData = env->GetByteArrayElements(rgba, JNI_FALSE);
    jbyte *yuvData = env->GetByteArrayElements(yuv, JNI_FALSE);
    int ret = func((const uint8_t *) rgbaData, rgba_stride, (uint8_t *) yuvData, y_stride,
                   (uint8_t *) (yuvData) + ySize, u_stride, (uint8_t *) (yuvData) + ySize + uSize,
                   v_stride, width, height);
    env->ReleaseByteArrayElements(rgba, rgbaData, JNI_OK);
    env->ReleaseByteArrayElements(yuv, yuvData, JNI_OK);
    return ret;
}

int
i420ToRgba(JNIEnv *env, jclass clazz, jbyteArray yuv, jint y_stride, jint u_stride, jint v_stride,
           jbyteArray rgba, jint rgba_stride, jint width, jint height,
           int (*func)(const uint8_t *, int, const uint8_t *, int, const uint8_t *, int, uint8_t *,
                       int, int, int)) {
    size_t ySize = (size_t) (y_stride * height);
    size_t uSize = (size_t) (u_stride * height >> 1);
    jbyte *rgbaData = env->GetByteArrayElements(rgba, JNI_FALSE);
    jbyte *yuvData = env->GetByteArrayElements(yuv, JNI_FALSE);
    int ret = func((const uint8_t *) yuvData, y_stride, (uint8_t *) yuvData + ySize, u_stride,
                   (uint8_t *) (yuvData) + ySize + uSize, v_stride, (uint8_t *) (rgbaData),
                   rgba_stride, width, height);
    env->ReleaseByteArrayElements(rgba, rgbaData, JNI_OK);
    env->ReleaseByteArrayElements(yuv, yuvData, JNI_OK);
    return ret;
}

int Jni_I420ToNV21(JNIEnv *env, jclass clazz, jbyteArray yuv420p, jbyteArray yuv420sp, jint width,
                   jint height, jboolean swapUV) {
    size_t ySize = (size_t) (width * height);
    size_t uSize = (size_t) (width * height >> 2);
    size_t stride[] = {0, uSize};
    jbyte *yuv420pData = env->GetByteArrayElements(yuv420p, JNI_FALSE);
    jbyte *yuv420spData = env->GetByteArrayElements(yuv420sp, JNI_FALSE);
    int ret = libyuv::I420ToNV21((const uint8_t *) yuv420pData, width,
                                 (const uint8_t *) (yuv420pData + ySize + stride[swapUV]),
                                 width >> 1,
                                 (const uint8_t *) (yuv420pData + ySize + stride[1 - swapUV]),
                                 width >> 1,
                                 (uint8_t *) yuv420spData, width,
                                 (uint8_t *) (yuv420spData + ySize), width, width, height);
    env->ReleaseByteArrayElements(yuv420p, yuv420pData, JNI_OK);
    env->ReleaseByteArrayElements(yuv420sp, yuv420spData, JNI_OK);
    return ret;
}

int Jni_NV21ToI420(JNIEnv *env, jclass clazz, jbyteArray yuv420p, jbyteArray yuv420sp, jint width,
                   jint height, jboolean swapUV) {
    size_t ySize = (size_t) (width * height);
    size_t uSize = (size_t) (width * height >> 2);
    size_t stride[] = {0, uSize};
    jbyte *yuv420pData = env->GetByteArrayElements(yuv420p, JNI_FALSE);
    jbyte *yuv420spData = env->GetByteArrayElements(yuv420sp, JNI_FALSE);
    int ret = libyuv::NV21ToI420((const uint8_t *) yuv420spData, width,
                                 (const uint8_t *) (yuv420spData + ySize), width,
                                 (uint8_t *) yuv420pData, width,
                                 (uint8_t *) (yuv420pData + ySize + stride[swapUV]), width >> 1,
                                 (uint8_t *) (yuv420pData + ySize + stride[1 - swapUV]), width >> 1,
                                 width, height);
    env->ReleaseByteArrayElements(yuv420p, yuv420pData, JNI_OK);
    env->ReleaseByteArrayElements(yuv420sp, yuv420spData, JNI_OK);
    return ret;
}

void
Jni_NV21Scale(JNIEnv *env, jclass clazz, jbyteArray src, jint width, jint height, jbyteArray dst,
              jint dst_width, jint dst_height, int mode) {
    size_t ySize = (size_t) (width * height);
    size_t dstYSize = (size_t) (dst_width * dst_height);
    jbyte *srcData = env->GetByteArrayElements(src, JNI_FALSE);
    jbyte *dstData = env->GetByteArrayElements(dst, JNI_FALSE);
    libyuv::ScalePlane((const uint8_t *) srcData, width, width, height,
                       (uint8_t *) dstData, dst_width, dst_width, dst_height,
                       (libyuv::FilterMode) mode);
    libyuv::ScalePlane((const uint8_t *) (srcData + ySize), width, width, height >> 1,
                       (uint8_t *) (dstData + dstYSize), dst_width, dst_width, dst_height >> 1,
                       (libyuv::FilterMode) mode);
    env->ReleaseByteArrayElements(src, srcData, JNI_OK);
    env->ReleaseByteArrayElements(dst, dstData, JNI_OK);
}

void
Jni_I420Scale(JNIEnv *env, jclass clazz, jbyteArray src, jint width, jint height, jbyteArray dst,
              jint dst_width, jint dst_height, int mode, jboolean swapUV) {
    int ySize = width * height;
    int swap[] = {0, ySize >> 2};
    size_t dstYSize = (size_t) (dst_width * dst_height);
    jbyte *srcData = env->GetByteArrayElements(src, JNI_FALSE);
    jbyte *dstData = env->GetByteArrayElements(dst, JNI_FALSE);
    libyuv::I420Scale((const uint8_t *) srcData, width, (const uint8_t *) (srcData + ySize),
                      width >> 1,
                      (uint8_t *) (srcData + ySize + (ySize >> 2)), width >> 1, width, height,
                      (uint8_t *) dstData, dst_width,
                      (uint8_t *) (dstData + dstYSize + swap[swapUV]), dst_width >> 1,
                      (uint8_t *) (dstData + dstYSize + swap[1 - swapUV]), dst_width >> 1,
                      dst_width, dst_height, (libyuv::FilterMode) mode);
    env->ReleaseByteArrayElements(src, srcData, JNI_OK);
    env->ReleaseByteArrayElements(dst, dstData, JNI_OK);
}

void Jni_NV21ToI420Rotate(JNIEnv *env, jclass clazz, jbyteArray src, jint width, jint height,
                          jbyteArray dst, jint de, jboolean swapUV) {
    int dst_stride = width;
    if (de % 180 != 0) {
        dst_stride = height;
    }
    size_t ySize = (size_t) (width * height);
    size_t swap[] = {0, ySize >> 2};
    jbyte *srcData = env->GetByteArrayElements(src, JNI_FALSE);
    jbyte *dstData = env->GetByteArrayElements(dst, JNI_FALSE);
    libyuv::RotatePlane((const uint8_t *) srcData, width, (uint8_t *) dstData, dst_stride,
                        width, height, (libyuv::RotationMode) de);
    libyuv::SplitRotateUV((const uint8_t *) (srcData + ySize), width,
                          (uint8_t *) (dstData + ySize + swap[swapUV]), dst_stride >> 1,
                          (uint8_t *) (dstData + ySize + swap[1 - swapUV]), dst_stride >> 1,
                          width >> 1, height >> 1, (libyuv::RotationMode) de);

    env->ReleaseByteArrayElements(src, srcData, JNI_OK);
    env->ReleaseByteArrayElements(dst, dstData, JNI_OK);
}

int
Jni_RgbaToI420WithStride(JNIEnv *env, jclass clazz, jint type, jbyteArray rgba, jint rgba_stride,
                         jbyteArray yuv, jint y_stride, jint u_stride, jint v_stride,
                         jint width, jint height) {
    uint8_t cType = (uint8_t) (type & 0x0F);
    return rgbaToI420(env, clazz, rgba, rgba_stride, yuv, y_stride, u_stride, v_stride, width,
                      height, rgbaToI420Func[cType]);
}

int
Jni_RgbaToI420(JNIEnv *env, jclass clazz, jint type, jbyteArray rgba, jbyteArray yuv, jint width,
               jint height) {
    uint8_t cType = (uint8_t) (type & 0x0F);
    int rgba_stride = ((type & 0xF0) >> 4) * width;
    int y_stride = width;
    int u_stride = width >> 1;
    int v_stride = u_stride;
    return rgbaToI420(env, clazz, rgba, rgba_stride, yuv, y_stride, u_stride, v_stride, width,
                      height, rgbaToI420Func[cType]);
}

int Jni_I420ToRgbaWithStride(JNIEnv *env, jclass clazz, jint type, jbyteArray yuv, jint y_stride,
                             jint u_stride, jint v_stride,
                             jbyteArray rgba, jint rgba_stride,
                             jint width, jint height) {
    uint8_t cType = (uint8_t) (type & 0x0F);
    return i420ToRgba(env, clazz, yuv, y_stride, u_stride, v_stride, rgba, rgba_stride, width,
                      height, i420ToRgbaFunc[cType]);
}

int
Jni_I420ToRgba(JNIEnv *env, jclass clazz, jint type, jbyteArray yuv, jbyteArray rgba, jint width,
               jint height) {
    uint8_t cType = (uint8_t) (type & 0x0F);
    int rgba_stride = ((type & 0xF0) >> 4) * width;
    int y_stride = width;
    int u_stride = width >> 1;
    int v_stride = u_stride;
    return i420ToRgba(env, clazz, yuv, y_stride, u_stride, v_stride, rgba, rgba_stride, width,
                      height, i420ToRgbaFunc[cType]);
}


//libyuv中，rgba表示abgrabgrabgr这样的顺序写入文件，java使用的时候习惯rgba表示rgbargbargba写入文件
static JNINativeMethod g_methods[] = {
        {"ARGBToNv21",       "([B[BII)I",      (void *) Jni_ARGBToNv21},
        {"NV21ToABGR",       "([B[BII)I",      (void *) Jni_Nv21ToABGR},
        {"RgbaToI420",       "(I[BI[BIIIII)I", (void *) Jni_RgbaToI420WithStride},
        {"RgbaToI420",       "(I[B[BII)I",     (void *) Jni_RgbaToI420},

        {"I420ToRgba",       "(I[BIII[BIII)I", (void *) Jni_I420ToRgbaWithStride},
        {"I420ToRgba",       "(I[B[BII)I",     (void *) Jni_I420ToRgba},

        {"I420ToNV21",       "([B[BIIZ)I",     (void *) Jni_I420ToNV21},
        {"NV21ToI420",       "([B[BIIZ)I",     (void *) Jni_NV21ToI420},
        {"NV21Scale",        "([BII[BIII)V",   (void *) Jni_NV21Scale},
        {"I420Scale",        "([BII[BIIIZ)V",  (void *) Jni_I420Scale},
        {"NV21ToI420Rotate", "([BII[BIZ)V",    (void *) Jni_NV21ToI420Rotate},
};

JNIEXPORT jint
JNI_OnLoad(JavaVM *vm, void *reserved) {
    JNIEnv *env = nullptr;

    if (vm->GetEnv((void **) &env, JNI_VERSION_1_4) != JNI_OK) {
        return JNI_ERR;
    }
    assert(env != nullptr);
    jclass clazz = env->FindClass(CLASS_JAVA);
    env->
            RegisterNatives(clazz, g_methods,
                            (int) (sizeof(g_methods) / sizeof((g_methods)[0])));

    return
            JNI_VERSION_1_4;
}

JNIEXPORT void JNI_OnUnload(JavaVM *jvm, void *reserved) {
    JNIEnv *env = nullptr;

    if (jvm->GetEnv((void **) &env, JNI_VERSION_1_4) != JNI_OK) {
        return;
    }
    jclass clazz = env->FindClass(CLASS_JAVA);
    env->UnregisterNatives(clazz);
}


#ifdef __cplusplus
}
#endif





